package pcap

import (
	"bytes"
	"encoding/base64"
	"io"
	"os"
	"regexp"
	"strconv"
	"strings"

	"github.com/kin9-0rz/gopcap"
)

// tcp 流
type tcpStream struct {
	source            string
	dest              string
	sourcePort        uint16
	destinationPort   uint16
	url               string
	protocol          string   // 协议
	datas             [][]byte // 数据
	nums              []uint32 // 存放ack/seq号
	received_datas    [][]byte
	transferred_datas [][]byte
}

// Parser Pcap解析器
type Parser struct {
	HTTPRequests []HTTPRequest
	DNSInfos     []DNSInfo
	MailInfos    MailInfos
	TDatas       []TransmissionData //存放每次交互传输的数据，不包括已解析的数据，如HTTP、DNS、SMTP。
	tcpStreams   []*tcpStream
	IPs          []string // 存放pcap中，出现的所有IP。
}

// NewParser 初始化
func NewParser() *Parser {
	return &Parser{
		HTTPRequests: make([]HTTPRequest, 0),
		DNSInfos:     make([]DNSInfo, 0),
		TDatas:       make([]TransmissionData, 0),
		MailInfos:    newMailInfos(),
		tcpStreams:   make([]*tcpStream, 0),
		IPs:          make([]string, 0),
	}
}

// Parse 解析目标pcap文件
func (p *Parser) Parse(path string) error {
	src, err := os.Open(path)
	defer src.Close()
	if err != nil {
		return err
	}

	err = p.ParseReader(src)
	if err != nil {
		return err
	}
	return nil
}

// ParseBytes 解析目标pcap bytes 格式
func (p *Parser) ParseBytes(bs []byte) error {
	err := p.ParseReader(bytes.NewReader(bs))
	if err != nil {
		return err
	}
	return nil
}

// ParseReader 解析目标pcap reader
func (p *Parser) ParseReader(src io.Reader) error {
	parsed, err := gopcap.Parse(src)
	if err != nil {
		if len(parsed.Packets) == 0 {
			return err
		}
	}

	num := len(parsed.Packets)
	for i := 0; i < num; i++ {
		packet := parsed.Packets[i]
		data := packet.Data
		if data == nil {
			continue
		}
		frame := data.(*gopcap.EthernetFrame)
		p.generateStream(frame)
	}

	for _, ts := range p.tcpStreams {
		if ts.protocol == "HTTP" {
			p.addHTTPRequest(ts)
		} else if ts.destinationPort == 25 {
			p.processSMTP(ts)
		} else {
			p.processTS(ts)
		}
	}

	ipRE := regexp.MustCompile(`^\d+\.\d+\.\d+\.\d+$`)
	for _, dnsInfo := range p.DNSInfos {
		for _, answer := range dnsInfo.Answers {
			p.AddIP(ipRE.FindString(answer.Addr))
		}
	}
	return nil
}

// AddIP add ip
func (p *Parser) AddIP(ip string) {
	for _, item := range p.IPs {
		if item == ip {
			return
		}
	}

	p.IPs = append(p.IPs, ip)
}

func (p *Parser) processTS(ts *tcpStream) {
	datas := make([][]byte, len(ts.datas))
	for _, data := range ts.datas {
		if len(data) < 3 {
			continue
		}
		datas = append(datas, data)
	}

	p.addTransmissionData(ts.source, ts.sourcePort, ts.dest, ts.destinationPort, bytes.Join(datas, []byte("")))
}

func decodeBase64(bs []byte) string {
	txt, err := base64.StdEncoding.DecodeString(string(bs))
	if err != nil {
		return ""
	}
	return string(txt)
}

func (p *Parser) processSMTP(ts *tcpStream) {
	if len(ts.datas) == 0 {
		return
	}
	username := ""
	password := ""
	from := ""
	to := ""
	subject := ""
	content := ""

	datas := bytes.Join(ts.datas, []byte(""))
	usrRE := regexp.MustCompile(`334 \w+\s*([\w=]+)\s*334 \w+\s*([\w=]+)\s*`)
	fromRE := regexp.MustCompile(`(?i:from)\s*:\s*<?([\w@\.]+)>?`)
	toRE := regexp.MustCompile(`(?i:to)\s*:\s*<?([\w@\.]+)>?`)
	subjectRE := regexp.MustCompile(`(?i:subject): (.+)[\n\r]`)
	contentRE := regexp.MustCompile(`354 (?:[\S\s]*)Content-.*\s*([\S\s]+)\.\s*\d{3} `)

	match := usrRE.FindSubmatch(datas)
	if len(match) < 3 {
		return
	}
	username = decodeBase64(match[1])
	password = decodeBase64(match[2])

	match = fromRE.FindSubmatch(datas)
	if len(match) < 2 {
		goto NEXT
	}
	from = string(match[1])

	match = toRE.FindSubmatch(datas)
	if len(match) < 2 {
		goto NEXT
	}
	to = string(match[1])

	match = subjectRE.FindSubmatch(datas)
	if len(match) < 2 {
		goto NEXT
	}
	subject = string(match[1])

	match = contentRE.FindSubmatch(datas)
	if len(match) < 2 {
		goto NEXT
	}
	content = string(match[1])

NEXT:
	p.MailInfos.addAccount(username, password)
	if from != "" {
		p.MailInfos.addContent(from, to, subject, content)
	}
}

func isHTTPRequest(ts *tcpStream) bool {
	for _, data := range ts.datas {
		if bytes.HasPrefix(data, []byte("GET ")) {
			return true
		}
		if bytes.HasPrefix(data, []byte("POST ")) {
			return true
		}
		if bytes.HasPrefix(data, []byte("HEAD ")) {
			return true
		}
		if bytes.HasPrefix(data, []byte("PUT ")) {
			return true
		}
		if bytes.HasPrefix(data, []byte("DELETE ")) {
			return true
		}
		if bytes.HasPrefix(data, []byte("TRACE ")) {
			return true
		}
		if bytes.HasPrefix(data, []byte("CONNECT ")) {
			return true
		}
		if bytes.HasPrefix(data, []byte("OPTIONS ")) {
			return true
		}
	}
	return false
}

func (p *Parser) addHTTPRequest(ts *tcpStream) {
	hr := HTTPRequest{}
	hr.SetRequestData(string(bytes.Join(ts.transferred_datas, []byte(""))))
	hr.SetResponeData(string(bytes.Join(ts.received_datas, []byte(""))))

	// 请求报文结构
	// 请求方法 空格 统一资源标识符 空格 HTTP协议版本 回车符 换行符

	// ! 请求方法 空格 统一资源标识符 回车符 换行符
	// ! 统一资源标识符 回车符 换行符
	// ! ... 统一资源标识符很长的时候，会被拆分
	// ! 统一资源标识符 回车符 换行符
	// ! HTTP协议版本 回车符 换行符

	// 头部字段名 冒号 值 回车符 换行符
	// 回车符 换行符
	// 请求体

	var request_line strings.Builder
	for _, item := range ts.transferred_datas {
		parts := strings.Split(string(item), "\r\n")
		flag := 0
		url := "http://"
		host := ""
		path := ""

		for _, line := range parts {
			request_line.WriteString(line)
			if strings.Contains(line, "HTTP/") {
				arr := strings.Split(request_line.String(), " ")
				hr.SetMethod(arr[0])
				path = arr[1]
				flag++
			} else if strings.HasPrefix(line, "Host: ") {
				arr := strings.Split(line, " ")
				host = arr[1]
				flag++
			}

			if flag == 2 {
				hr.SetURL(url + host + path)
				p.HTTPRequests = append(p.HTTPRequests, hr)
				flag = 0
				host = ""
				path = ""
			}

			request_line.Reset()
		}

		// if flag == 2 {
		// 	hr.SetURL(url + host + path)
		// 	break
		// }
	}
	p.HTTPRequests = append(p.HTTPRequests, hr)
}

func (p *Parser) generateStreamGo(frame *gopcap.EthernetFrame, ch chan struct{}) {
	linkData := frame.LinkData()
	if linkData == nil {
		ch <- struct{}{}
		return
	}

	// 解析传输数据
	if frame.EtherType == gopcap.ETHERTYPE_IPV4 {
		pkt, ok := linkData.(*gopcap.IPv4Packet)
		if !ok {
			ch <- struct{}{}
			return
		}
		// 传输层主要有TCP和UDP两种协议
		if pkt.Protocol == gopcap.IPP_TCP {
			p.generateTCPStream(pkt)
		} else if pkt.Protocol == gopcap.IPP_UDP {
			p.parseUDPDatagram(pkt)
		} else {
			// TODO 其他协议暂不处理
		}
	} else if frame.EtherType == gopcap.ETHERTYPE_IPV6 {
		// TODO 暂时不处理IPV6
	}

	ch <- struct{}{}
}

func (p *Parser) generateStream(frame *gopcap.EthernetFrame) {
	linkData := frame.LinkData()
	if linkData == nil {
		return
	}

	// 解析传输数据
	if frame.EtherType == gopcap.ETHERTYPE_IPV4 {
		pkt, ok := linkData.(*gopcap.IPv4Packet)
		if !ok {
			return
		}
		// 传输层主要有TCP和UDP两种协议
		if pkt.Protocol == gopcap.IPP_TCP {
			p.generateTCPStream(pkt)
		} else if pkt.Protocol == gopcap.IPP_UDP {
			p.parseUDPDatagram(pkt)
		} else {
			// TODO 其他协议暂不处理
		}
	} else if frame.EtherType == gopcap.ETHERTYPE_IPV6 {
		// TODO 暂时不处理IPV6
	}
}

func isHTTP(data []byte) bool {
	if bytes.HasPrefix(data, []byte("GET ")) {
		return true
	}
	if bytes.HasPrefix(data, []byte("POST ")) {
		return true
	}
	if bytes.HasPrefix(data, []byte("HEAD ")) {
		return true
	}
	if bytes.HasPrefix(data, []byte("PUT ")) {
		return true
	}
	if bytes.HasPrefix(data, []byte("DELETE ")) {
		return true
	}
	if bytes.HasPrefix(data, []byte("TRACE ")) {
		return true
	}
	if bytes.HasPrefix(data, []byte("CONNECT ")) {
		return true
	}
	if bytes.HasPrefix(data, []byte("OPTIONS ")) {
		return true
	}
	return false
}

func (p *Parser) generateTCPStream(pkt *gopcap.IPv4Packet) {
	data := pkt.InternetData()
	sourceAddress := ConvertIP(pkt.SourceAddress)
	destAddress := ConvertIP(pkt.DestAddress)

	p.AddIP(sourceAddress)
	p.AddIP(destAddress)

	tcpSegment, ok := data.(*gopcap.TCPSegment)
	if !ok {
		return
	}

	ack := tcpSegment.AckNumber
	seq := tcpSegment.SequenceNumber

	// SYN ack=0 seq=随机数，tcp.nxtseq=1 ack和nxtseq是默认值
	if ack == 0 {
		ts := &tcpStream{
			source:            sourceAddress,
			dest:              destAddress,
			sourcePort:        tcpSegment.SourcePort,
			destinationPort:   tcpSegment.DestinationPort,
			protocol:          "tcp",
			nums:              []uint32{seq},
			datas:             make([][]byte, 0),
			transferred_datas: make([][]byte, 0),
			received_datas:    make([][]byte, 0),
		}
		p.tcpStreams = append(p.tcpStreams, ts)
		return
	}

	for _, ts := range p.tcpStreams {
		for _, num := range ts.nums {
			if num == ack {
				ts.nums = append(ts.nums, seq)
			} else if num == ack-1 {
				ts.nums = append(ts.nums, ack)
				ts.nums = append(ts.nums, seq)
			} else if num == seq {
				ts.nums = append(ts.nums, ack)
			} else {
				continue
			}

			data := tcpSegment.TransportData()
			if len(data) == 0 {
				return
			}
			ts.datas = append(ts.datas, data)

			if ts.source == sourceAddress {
				ts.transferred_datas = append(ts.transferred_datas, data)
				if isHTTP(data) {
					ts.protocol = "HTTP"
				}
			} else {
				ts.received_datas = append(ts.received_datas, data)
			}
			return
		}
	}
}

func (p *Parser) parseUDPDatagram(pkt *gopcap.IPv4Packet) {
	data := pkt.InternetData()

	datagram, ok := data.(*gopcap.UDPDatagram)
	if !ok {
		return
	}

	// 解析DNS数据
	if datagram.SourcePort == 53 {
		tData := datagram.TransportData()
		// 如果源端口是53，则可能是DNS协议
		dns := DNS{}
		dns.FromBytes(tData)
		p.AddDNSInfo(dns.dnsInfo)
	}
	// 其他UDP协议暂不解析
}

// TODO 处理UDP
func (p *Parser) generateUDPStream(pkt *gopcap.IPv4Packet) {
	// data := pkt.InternetData()
	sourceAddress := ConvertIP(pkt.SourceAddress)
	destAddress := ConvertIP(pkt.DestAddress)

	p.AddIP(sourceAddress)
	p.AddIP(destAddress)

	// udpDatagram, ok := data.(*gopcap.UDPDatagram)
	// if !ok {
	// 	return
	// }
}

// MailInfos 邮箱信息
type MailInfos struct {
	Accounts []MailAccount `json:"accounts,omitempty"`
	Contents []MailContent `json:"contents,omitempty"`
}

func newMailInfos() MailInfos {
	return MailInfos{
		Accounts: make([]MailAccount, 0),
		Contents: make([]MailContent, 0),
	}
}

// MailAccount 邮箱账号
type MailAccount struct {
	User string `json:"user,omitempty"`
	Pwd  string `json:"pwd,omitempty"`
}

// MailContent 邮件内容
type MailContent struct {
	From    string `json:"from,omitempty"`
	To      string `json:"to,omitempty"`
	Subject string `json:"subject,omitempty"`
	Content string `json:"content,omitempty"`
}

func (m *MailInfos) addAccount(user string, pwd string) {
	ma := MailAccount{
		User: user,
		Pwd:  pwd,
	}
	for _, a := range m.Accounts {
		if a == ma {
			return
		}
	}
	m.Accounts = append([]MailAccount{ma}, m.Accounts...)
}

func (m *MailInfos) addContent(from string, to string, subject string, content string) {
	mc := MailContent{
		From:    from,
		To:      to,
		Subject: subject,
		Content: content,
	}
	for _, i := range m.Contents {
		if i == mc {
			return
		}
	}
	m.Contents = append([]MailContent{mc}, m.Contents...)
}

// TransmissionData 存放传输的数据，TCP/UDP等，未解析的数据
type TransmissionData struct {
	Source          string // 来源的IP/域名/URL
	SourcePort      uint16 // 来源的IP/域名/URL
	Destination     string // 目的地的IP/域名/URL
	DestinationPort uint16 // 端口
	Data            []byte //传输的数据
	// TODO 是否需要加上协议？
}

func (p *Parser) addTransmissionData(
	source string,
	sPort uint16,
	dest string,
	dPort uint16,
	data []byte) {
	td := TransmissionData{
		Source:          source,
		SourcePort:      sPort,
		Destination:     dest,
		DestinationPort: dPort,
		Data:            data,
	}
	p.TDatas = append([]TransmissionData{td}, p.TDatas...)
}

// ConvertIP 将bytes格式的IP，转为字符串
func ConvertIP(bytes []byte) string {
	var str = ""
	for i := 0; i < len(bytes); i++ {
		str += strconv.Itoa(int(bytes[i]))
		if i < 3 {
			str += "."
		}
	}
	return str
}

// AddDNSInfo add dns info
func (p *Parser) AddDNSInfo(d DNSInfo) {
	domain := d.Domain
	for _, item := range p.DNSInfos {
		if item.Domain == domain {
			return
		}
	}
	p.DNSInfos = append(p.DNSInfos, d)
}
